# -*- coding: utf-8 -*-

'''
Created on 2017年12月3日
@author: Nick
'''

'''
   获取并可视化股票数据
'''
import pandas as pd
import pandas_datareader.data as web   # 导入 Data 需要导入的包和模块
import datetime
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter
from matplotlib.dates import WeekdayLocator
from matplotlib.dates import DayLocator
from matplotlib.dates import MONDAY
from matplotlib.dates import date2num
from matplotlib.finance import candlestick_ohlc
 
def pandas_candlestick_ohlc(dat, stick = "day", otherseries = None):
    """
    :param dat: pandas DataFrame object with datetime64 index, 
    and float columns "Open", "High", "Low", and "Close", 
    likely created via DataReader from "yahoo"
    
    :param stick: A string or number indicating the period of time 
    covered by a single candlestick. Valid string inputs 
    include "day", "week", "month", and "year", ("day" default), a
    nd any numeric input indicates the number of trading days included in a period
    
    :param otherseries: An iterable that will be coerced into a list, 
    containing the columns of dat that hold other series to be plotted as lines
    """
    mondays = WeekdayLocator(MONDAY)    # major ticks on the mondays
    alldays = DayLocator()              # minor ticks on the days
    # dayFormatter = DateFormatter('%d')     # e.g., 12
 
    # Create a new DataFrame which includes OHLC data 
    # for each period specified by stick input
    transdat = dat.loc[:,["Open", "High", "Low", "Close"]]
    if (type(stick) == str):
        if stick == "day":
            plotdat = transdat
            stick = 1 # Used for plotting
        elif stick in ["week", "month", "year"]:
            if stick == "week":
                transdat["week"] = pd.to_datetime(transdat.index).map(lambda x: x.isocalendar()[1])  # Identify weeks
            elif stick == "month":
                transdat["month"] = pd.to_datetime(transdat.index).map(lambda x: x.month)  # Identify months
            transdat["year"] = pd.to_datetime(transdat.index).map(lambda x: x.isocalendar()[0])  # Identify years
            # Group by year and other appropriate variable
            grouped = transdat.groupby(list(set(["year",stick])))  
            # Create empty data frame containing what will be plotted
            plotdat = pd.DataFrame({"Open": [], "High": [], "Low": [], "Close": []}) 
            for group in grouped:
                plotdat = plotdat.append(pd.DataFrame({"Open": group.iloc[0,0],
                                    "High": max(group.High),
                                    "Low": min(group.Low),
                                    "Close": group.iloc[-1,3]},
                                    index = [group.index[0]]))
            if stick == "week": stick = 5
            elif stick == "month": stick = 30
            elif stick == "year": stick = 365
 
    elif (type(stick) == int and stick >= 1):
        transdat["stick"] = [np.floor(i / stick) for i in range(len(transdat.index))]
        grouped = transdat.groupby("stick")
        plotdat = pd.DataFrame({"Open": [], "High": [], "Low": [], "Close": []}) # Create empty data frame containing what will be plotted
        for group in grouped:
            plotdat = plotdat.append(pd.DataFrame({"Open": group.iloc[0,0],
                                        "High": max(group.High),
                                        "Low": min(group.Low),
                                        "Close": group.iloc[-1,3]},
                                       index = [group.index[0]]))
 
    else:
        raise ValueError('Valid inputs to argument "stick" include the strings "day", "week", "month", "year", or a positive integer')
 
 
    # Set plot parameters, including the axis object ax used for plotting
    fig, ax = plt.subplots()
    fig.subplots_adjust(bottom=0.2)
    if plotdat.index[-1] - plotdat.index[0] < pd.Timedelta('730 days'):
        weekFormatter = DateFormatter('%b %d')  # e.g., Jan 12
        ax.xaxis.set_major_locator(mondays)
        ax.xaxis.set_minor_locator(alldays)
    else:
        weekFormatter = DateFormatter('%b %d, %Y')
    ax.xaxis.set_major_formatter(weekFormatter)
 
    ax.grid(True)
 
    # Create the candelstick chart
    candlestick_ohlc(ax, list(zip(list(date2num(plotdat.index.tolist())), plotdat["Open"].tolist(), plotdat["High"].tolist(),
                      plotdat["Low"].tolist(), plotdat["Close"].tolist())),
                      colorup = "black", colordown = "red", width = stick * .4)
 
    # Plot other series (such as moving averages) as lines
    if otherseries != None:
        if type(otherseries) != list:
            otherseries = [otherseries]
        dat.loc[:,otherseries].plot(ax = ax, lw = 1.3, grid = True)
 
    ax.xaxis_date()
    ax.autoscale_view()
    plt.setp(plt.gca().get_xticklabels(), rotation=45, horizontalalignment='right')
 
    plt.show()

start = datetime.datetime(2017,1,1)
end = datetime.datetime(2017,12,1)
apple = web.DataReader("AAPL", "yahoo", start, end) 

pandas_candlestick_ohlc(apple)
